/*
 * Copyright 2004-2024 H2 Group. Multiple-Licensed under the MPL 2.0,
 * and the EPL 1.0 (https://h2database.com/html/license.html).
 * Initial Developer: H2 Group
 */
package org.h2.expression.function;

import com.softdev.system.generator.controller.GeneratorController;
import org.h2.command.Parser;
import org.h2.engine.SessionLocal;
import org.h2.expression.Expression;
import org.h2.expression.TypedValueExpression;
import org.h2.message.DbException;
import org.h2.value.TypeInfo;
import org.h2.value.Value;
import org.h2.value.ValueArray;
import org.h2.value.ValueDouble;
import org.h2.value.ValueVarchar;

import java.nio.charset.StandardCharsets;
import java.util.Base64;
import java.util.concurrent.ConcurrentHashMap;

/**
 * An string function with multiple arguments.
 */
public final class LLMFunction extends FunctionN {

    private static ConcurrentHashMap<String, Integer> NAME2INDEX;
    private static final String[] NAMES = { //
            "LLMGENCODE",
            "BASE64ENCODE",
            "BASE64DECODE",
            "VCT_COSINE",
            "VCT_ADJUST",

            //EXCEL生成
            "ACT_CHATLATS",
            "ACT_CHATMAP",
            "ACT_CHATPLAN",
            "ACT_CHATRELATE",
            "ACT_CHATREREVS",
            "ACT_CHATWT",
            "ACT_LOADMAP",
            "ACT_LOADPLAN",
            "ACT_LOADRELATE",
            "ACT_LOADREVS",
            "ACT_LOADWT",
            "APIACT",
            "APIBASE",
            "APIEMBED",
            "APIGRP",
            "APIMODELS",
            "APIRFT",
            "APISERVER",
            "APISFT",
            "APISQLVECTOR",
            "APITOL",
            "APITOOLS",
            "APIVANNA",
            "APIVCTSVR",
            "APIVECTOR",
            "DAT_CREATE",
            "DAT_DEL",
            "DAT_READ",
            "DAT_UPDATE",
            "EPT_LATS",
            "FILE_PDF",
            "FILE_READALL",
            "GRP_AGENT",
            "GRP_CHAT",
            "GRP_FIND",
            "GRP_LATS",
            "GRP_LOADGRP",
            "GRP_LOADJSON",
            "GRP_QUERY",
            "GRP_REGHPL",
            "GRP_REGJSN",
            "GRP_REGREC",
            "GRP_THINK",
            "INFO",
            "LLMAGENT",
            "LLMCHAT",
            "LLMINVOKE",
            "LLMPIP",
            "ML_APRIORI",
            "ML_FORESTCLS",
            "ML_FORESTREG",
            "ML_FPGROWTH",
            "ML_KMENS",
            "ML_KNNCLS",
            "ML_KNNREG",
            "ML_LINEARCLS",
            "ML_LINEARREG",
            "ML_LOGISTIC",
            "ML_SVMCLS",
            "ML_SVMREG",
            "ML_TREECLS",
            "ML_TREEREG",
            "NLP_CLS",
            "NLP_CORRECTION",
            "NLP_ENTITY",
            "NLP_ENTITYEXTRACT",
            "NLP_FORMAT",
            "NLP_INFER",
            "NLP_MAP",
            "NLP_MASK",
            "NLP_PIC",
            "NLP_REDUCE",
            "NLP_RELATION",
            "NLP_SEGMENT",
            "NLP_SENTIMENT",
            "NLP_SIMILARITY",
            "NLP_SUMMARY",
            "NLP_TRACT",
            "NLP_TRANSLATE",
            "NLP_VOICE",
            "PIC_BRIA",
            "PIC_CLS",
            "PIC_CUT",
            "PIC_OCR",
            "PIC_OFA",
            "PIC_RECOGNITION",
            "PIC_SEARCH",
            "PIC_VOICETXT",
            "PMP_ADJUST",
            "PROMPINIT",
            "RAG_CHAT",
            "RAG_CHATGB",
            "RAG_CHATNODE",
            "RAG_CHATRIP",
            "RAG_CORRECT",
            "RAG_DOC2QA",
            "RAG_FIXQA",
            "RAG_GRAPH",
            "RAG_LOADFILE",
            "RAG_LOADGB",
            "RAG_LOADNODE",
            "RAG_LOADQA",
            "RAG_LOADTXT",
            "RAG_NAME",
            "RAG_NEO4J",
            "RAG_QUERY",
            "RAG_SQL",
            "RAG_WEBSEEK",
            "RANK_BM25",
            "RFT_GEVAL",
            "RFT_LOADTXT",
            "RFT_QUERY",
            "RFT_RELEVANCY",
            "SAVEENV",
            "SFT_MASK",
            "SFT_RAG",
            "TEST_LLM",
            "TOL_CHAT",
            "TOL_LOADCODE",
            "TOL_LOADFILE",
            "TOL_LOADFUNC",
            "TOL_MATH",
            "TOL_QUERY",
            "TOL_WEB",
            "TOL_WEB_SEEK",
            "TOL_WRIGHT",
    };

    private final int function;

    public LLMFunction(Expression arg1, int function) {
        super(new Expression[]{arg1});
        this.function = function;
    }

    public LLMFunction(Expression arg1, Expression arg2, int function) {
        super(arg2 == null ? new Expression[]{arg1} : new Expression[]{arg1, arg2});
        this.function = function;
    }

    public LLMFunction(Expression arg1, Expression arg2, Expression arg3, int function) {
        super(arg3 == null ? new Expression[]{arg1, arg2} : new Expression[]{arg1, arg2, arg3});
        this.function = function;
    }

    public LLMFunction(Expression arg1, Expression arg2, Expression arg3, Expression arg4, int function) {
        super(arg4 == null ? new Expression[]{arg1, arg2, arg3} : new Expression[]{arg1, arg2, arg3, arg4});
        this.function = function;
    }

    public LLMFunction(Expression arg1, Expression arg2, Expression arg3, Expression arg4, Expression arg5, int function) {
        super(arg5 == null ? new Expression[]{arg1, arg2, arg3, arg4} : new Expression[]{arg1, arg2, arg3, arg4, arg5});
        this.function = function;
    }

    @Override
    public Value getValue(SessionLocal session) {
        String tmpVal;
        Value v;
        String pyName;
        String val, val1, val2;
        String path = session.getDatabase().getDatabasePath();
        String dbDir = "";
        if (path != null) {
            int start = path.lastIndexOf("/");
            dbDir = path.substring(0, start);
        }
        double cosine;
        String funcName = NAMES[function];
        switch (funcName) {
            case "VCT_COSINE":
                val1 = args[0].getValue(session).getString();
                val2 = args[1].getValue(session).getString();
                cosine = ChatgptFunction.calculateCosineSimilarity(val1, val2);
                return ValueDouble.get(cosine);
            case "VCT_ADJUST":
                val1 = args[0].getValue(session).getString();
                val2 = args[1].getValue(session).getString();
                cosine = ChatgptFunction.calculateCosineSimilarity_adjust(val1, val2);
                return ValueDouble.get(cosine);
            case "BASE64ENCODE":
                // 编码成Base64
                v = args[0].getValue(session);
                val = v.getString();
                val1 = Base64.getEncoder().encodeToString(val.getBytes(StandardCharsets.UTF_8));
                //System.out.println("Base64 Encoded: " + encodedString);
                return ValueVarchar.get(val1, session);
            case "BASE64DECODE":
                // Base64解码
                v = args[0].getValue(session);
                val = v.getString();
                byte[] decodedBytes = Base64.getDecoder().decode(val);
                val1 = new String(decodedBytes, StandardCharsets.UTF_8);
                //System.out.println("Base64 Encoded: " + encodedString);
                return ValueVarchar.get(val1, session);
            case "LLMGENCODE":
                //pyName = "llmtxtcode";
                v = args[0].getValue(session);
                val = v.getString();
                v = args[1].getValue(session);
                val1 = v.getString();
                GeneratorController gen = new GeneratorController();
                try {
                    val2 = gen.generateCode(dbDir, val, val1);
                } catch (Exception e) {
                    val2 = e.getMessage();
                }
                return ValueVarchar.get(val2, session);
            default:
                pyName = NAMES[function].toLowerCase();
                System.out.println("[bzm]LLMFunction pyname=" + pyName);
                break;
        }
        tmpVal = ChatgptFunction.runPython(pyName, args, session, 0);
        return ValueVarchar.get(tmpVal, session);
    }

    @Override
    public Expression optimize(SessionLocal session) {
        boolean allConst = optimizeArguments(session, true);
        type = TypeInfo.TYPE_VARCHAR;
        if (allConst) {
            return TypedValueExpression.getTypedIfNull(getValue(session), type);
        }
        return this;
    }

    @Override
    public String getName() {
        return NAMES[function];
    }

    /**
     * 根据函数名称 获取NAMES的编号
     *
     * @param name 函数名称
     * @return 返回NAMES的位置
     */
    public static Integer getFunctionIndex(String name) {
        if (LLMFunction.NAME2INDEX == null) {
            LLMFunction.NAME2INDEX = new ConcurrentHashMap<>();
            for (int i = 0; i < LLMFunction.NAMES.length; i++) {
                LLMFunction.NAME2INDEX.put(LLMFunction.NAMES[i], i);
            }
        }
        return LLMFunction.NAME2INDEX.get(name);
    }

    /**
     * 处理函数
     *
     * @param funcName 函数名称
     * @param parser   获取参数
     * @return 返回值
     */
    public static LLMFunction ProcessFunction(String funcName, Parser parser) {
        Integer index = LLMFunction.getFunctionIndex(funcName);
        if (index == null) {
            return null;
        }
        switch (funcName) {
            case "LLMGENCODE":
                return new LLMFunction(parser.readExpression(), parser.readNextArgument(), parser.readIfArgument(), index);
            //2个参数
            case "BASE64ENCODE":
            case "BASE64DECODE":
            case "VCT_COSINE":
            case "VCT_ADJUST":
                return new LLMFunction(parser.readExpression(), parser.readIfArgument(), index);

            //1个参数
            case "GRP_FIND":
            case "RAG_NAME":
            case "FILE_PDF":
            case "GRP_REDUCE":
            case "NLP_CORRECTION":
            case "NLP_ENTITY":
            case "NLP_RELATION":
            case "NLP_TRACT":
                return new LLMFunction(parser.readSingleArgument(), index);
            //4个参数
            //case "RAG_CHAT":
            //    return new LLMFunction(parser.readExpression(), parser.readNextArgument(), parser.readNextArgument(), parser.readIfArgument(), index);
            //5个参数
            case "ACT_LOADMAP":
            case "ACT_LOADPLAN":
            case "ACT_LOADRELATE":
            case "ACT_LOADREVS":
            case "ACT_LOADWT":
            case "RAG_QUERY":
            case "RAG_LOADQA":
            case "RAG_CHAT":
            case "RFT_GEVAL":
            case "RFT_RELEVANCY":
            case "TOL_LOADFUNC":
                return new LLMFunction(parser.readExpression(), parser.readNextArgument(), parser.readNextArgument(), parser.readNextArgument(), parser.readIfArgument(), index);
            default:
                //3个参数
                //case "RAG_QUERY":
                return new LLMFunction(parser.readExpression(), parser.readNextArgument(), parser.readIfArgument(), index);
        }
    }
}
